## 获取voc标签中所有的类别及其对应的个数
import os
import xml.etree.ElementTree as ET
from collections import defaultdict

def count_categories_and_images(file_path, category_count):
    # 解析 XML 文件
    tree = ET.parse(file_path)
    root = tree.getroot()

    # 获取图片名称
    filename_tag = root.find('filename')
    if filename_tag is not None:
        filename = filename_tag.text

        # 遍历所有的 <object> 标签并统计 <name> 标签
        for obj in root.iter('object'):
            name_tag = obj.find('name')
            if name_tag is not None:
                category = name_tag.text
                # 统计该类别数量并记录对应的图片文件名
                category_count[category]['count'] += 1
                category_count[category]['images'].append(filename)

def count_categories_in_folder(folder_path):
    category_count = defaultdict(lambda: {'count': 0, 'images': []})  # 存储每个类别的计数和图片列表
    # 遍历文件夹中的所有 XML 文件
    for filename in os.listdir(folder_path):
        if filename.endswith('.xml'):
            file_path = os.path.join(folder_path, filename)
            count_categories_and_images(file_path, category_count)
            # print(f"Processed: {filename}")
    
    # 返回统计结果
    return category_count

# 设置标签文件夹路径
folder_path = '/home/shuai/Documents/mydataset/my_usv_aug/AUG-Seaships/Annotations'  # 替换为你实际的标签文件夹路径

# 执行统计操作
category_count = count_categories_in_folder(folder_path)

# 输出统计结果
print("\nCategory count:")
categories = []  # 存储所有类别
for category, data in category_count.items():
    print(f"{category}: {data['count']} images")
    # 看是否需要打印每一个类别对应的图片是哪个
    # print(f"Images: {', '.join(data['images'])}")
    # 将类别名添加到列表中
    categories.append(category)

# 打印所有类别列表
print("\nAll categories:", categories)